package com.fmy.test.config.dynamic.algorithm;

import com.fmy.test.config.dynamic.enums.DBType;
import org.springframework.util.Assert;

import java.util.Map;
import java.util.Set;
import java.util.TreeMap;

/**
 * <p>
 * 权重算法
 * <p>
 *
 * @ClassName WeighAlgorithm
 * @Aulthor Mr.fmy
 * @Version 1.0
 * @Date 2019/8/26 12:56
 **/
public class WeighAlgorithm implements DBAlgorithm {

    /**
     * key 为 数据源bean 名称 value 为权重比
     */
    private Map<String, Integer> writeMap; //写实例权重

    private Map<String, Integer> readMap; //读实例权重

    private final TreeMap<Integer, String> WRITE_TREE_MAP, READ_TREE_MAP;

    private Integer write_weight_count = 0, read_weight_count = 0;

    public WeighAlgorithm() {
        WRITE_TREE_MAP = new TreeMap<>();
        READ_TREE_MAP = new TreeMap<>();
    }

    @Override
    public String getDB(DBType dbType) {
        int total = 0;
        TreeMap<Integer, String> treeMap = null;
        if (DBType.WRITE_DB.equals(dbType)) {
            treeMap = WRITE_TREE_MAP;//写
            total = write_weight_count;
        } else if (DBType.READ_DB.equals(dbType)) {
            treeMap = READ_TREE_MAP;//读
            total = read_weight_count;
        }
        //生产随机数
        int random = (int) (Math.random() * total);
        assert treeMap != null;
        return treeMap.get(treeMap.ceilingKey(random));
    }

    public synchronized void setWriteMap(Map<String, Integer> writeMap) {
        Assert.notNull(writeMap, "请配置`写`数据源的权重");
        this.writeMap = writeMap;//初始化写
        init(DBType.WRITE_DB);
    }

    public synchronized void setReadMap(Map<String, Integer> readMap) {
        Assert.notNull(readMap, "请配置`读`数据源的权重");
        this.readMap = readMap;
        //初始化读
        init(DBType.READ_DB);
    }

    private void init(DBType dbType) {
        TreeMap<Integer, String> treeMap = null;
        Map<String, Integer> hashMap = null;
        if (DBType.WRITE_DB.equals(dbType)) {
            treeMap = WRITE_TREE_MAP;//写
            hashMap = writeMap;
        } else if (DBType.READ_DB.equals(dbType)) {
            treeMap = READ_TREE_MAP;//读
            hashMap = readMap;
        }
        assert treeMap != null;
        treeMap.clear();
        int weightCount = 0;
        Set<String> setMap = hashMap.keySet();
        for (String key : setMap) {
            int weight = hashMap.get(key);//获取权重
            for (int i = 0; i < weight; i++) {
                treeMap.put(weightCount + i, key);
            }
            weightCount += weight;
        }
        if (DBType.WRITE_DB.equals(dbType)) {
            write_weight_count = weightCount;
        } else if (DBType.READ_DB.equals(dbType)) {
            read_weight_count = weightCount;
        }
    }
}
